#include "downloader.hh"

#ifdef _MSC_VER
#define CURL_STATICLIB
#endif

#include <curl/curl.h>

#include <fstream>
#include <sstream>

struct DownloadTask {
    ~DownloadTask() {
        delete ofile;
    }
    Downloader *host = nullptr;
    int id = 0;
    CURL *handle = nullptr;
    std::ostream *ofile = nullptr;
    uint64_t filetime = 0ULL;
    uint64_t curr = 0ULL;
    uint64_t total = 0ULL;
    Downloader::State state = Downloader::Paused;
    std::string url;
    void (*finishCallback)(DownloadTask*) = nullptr;
    void(*progressCallback)(DownloadTask*, uint64_t, uint64_t) = nullptr;
    void *data = nullptr;
    uintptr_t param = 0;

    bool noBody = false;
};

Downloader::Downloader() {
    curl_global_init(CURL_GLOBAL_DEFAULT);
    handle_ = curl_multi_init();
}

Downloader::~Downloader() {
    for (auto &p: tasks_) {
        if (p.second->state == Paused || p.second->state == Downloading)
            curl_multi_remove_handle(handle_, p.second->handle);
        curl_easy_cleanup(p.second->handle);
        delete p.second;
    }
    tasks_.clear();
    curl_multi_cleanup(handle_);
    curl_global_cleanup();
}

void Downloader::process() {
    if (running_) {
        int r;
        curl_multi_perform(handle_, &r);
        if (r != running_) {
            int msg_in_queue;
            running_ = r;
            do {
                CURLMsg *msg = curl_multi_info_read(handle_, &msg_in_queue);
                if (msg == nullptr) break;
                if (msg->msg == CURLMSG_DONE) {
                    DownloadTask *dt;
                    curl_easy_getinfo(msg->easy_handle, CURLINFO_PRIVATE, &dt);
                    if (dt != nullptr) {
                        if (msg->data.result == CURLE_OK) {
                            curl_easy_getinfo(dt->handle, CURLINFO_FILETIME_T, &dt->filetime);
                            if (dt->total == 0ULL)
                                curl_easy_getinfo(dt->handle, CURLINFO_CONTENT_LENGTH_DOWNLOAD_T, &dt->total);
                            if (dt->ofile) dt->ofile->flush();
                            curl_multi_remove_handle(handle_, dt->handle);
                            curl_easy_cleanup(dt->handle);
                            dt->state = Finished;
                            if (dt->finishCallback) dt->finishCallback(dt);
                            tasks_.erase(dt->id);
                            delete dt;
                        } else if (msg->data.result == CURLE_PARTIAL_FILE || msg->data.result == CURLE_RECV_ERROR) {
                            curl_multi_remove_handle(handle_, dt->handle);
                            curl_easy_cleanup(dt->handle);
                            dt->handle = curl_easy_init();
                            addCommon(dt, false, dt->ofile->tellp());
                            curl_multi_perform(handle_, &running_);
                        }
                    }
                }
            } while (msg_in_queue);
        }
    }
}

DownloadTask *Downloader::addUrl(const std::string &url, bool paused /*= false*/) {
    DownloadTask *dt = addCommon(url, paused);
    if (dt == nullptr) {
        return nullptr;
    }
    curl_easy_setopt(dt->handle, CURLOPT_WRITEFUNCTION, _writecb);
    curl_easy_setopt(dt->handle, CURLOPT_WRITEDATA, dt);
    dt->ofile = new std::ostringstream;
    curl_multi_perform(handle_, &running_);
    return dt;
}

DownloadTask *Downloader::addUrlFile(const std::string &url, const std::string &localfile, bool paused) {
    auto ofile = new std::ofstream;
    ofile->open(localfile, std::ios::out | std::ios::binary);
    if (!ofile->is_open()) {
        delete ofile;
        return nullptr;
    }
    DownloadTask *dt = addCommon(url, paused);
    if (dt == nullptr) {
        delete ofile;
        return nullptr;
    }
    curl_easy_setopt(dt->handle, CURLOPT_WRITEFUNCTION, _writecb);
    curl_easy_setopt(dt->handle, CURLOPT_WRITEDATA, dt);
    dt->ofile = ofile;
    curl_multi_perform(handle_, &running_);
    return dt;
}

DownloadTask *Downloader::addUrlHead(const std::string &url) {
    DownloadTask *dt = addCommon(url);
    if (dt == nullptr) {
        return nullptr;
    }
    dt->noBody = true;
    curl_easy_setopt(dt->handle, CURLOPT_NOBODY, 1);
    curl_multi_perform(handle_, &running_);
    return dt;
}

bool Downloader::resume(DownloadTask *dt) {
    if (dt->state != Paused) return false;
    curl_easy_pause(dt->handle, CURLPAUSE_CONT);
    dt->state = Downloading;
    return true;
}

bool Downloader::pause(DownloadTask *dt) {
    if (dt->state != Downloading) return false;
    curl_easy_pause(dt->handle, CURLPAUSE_ALL);
    dt->state = Paused;
    return true;
}

bool Downloader::cancel(DownloadTask *dt) {
    curl_easy_reset(dt->handle);
    dt->state = Cancelled;
    return true;
}

Downloader::State Downloader::state(DownloadTask *dt) {
    return dt->state;
}

void *Downloader::getUserData(DownloadTask *dt) {
    return dt->data;
}

uintptr_t Downloader::getUserParam(DownloadTask *dt) {
    return dt->param;
}

void Downloader::setUserData(DownloadTask *dt, void *data) {
    dt->data = data;
}

void Downloader::setUserParam(DownloadTask * dt, uintptr_t param) {
    dt->param = param;
}

void Downloader::setFinishCallback(DownloadTask *dt, void(*finishCallback)(DownloadTask*)) {
    dt->finishCallback = finishCallback;
}

void Downloader::setProgressCallback(DownloadTask *dt, void(*progressCallback)(DownloadTask*, uint64_t, uint64_t)) {
    if (progressCallback) {
        curl_easy_setopt(dt->handle, CURLOPT_XFERINFOFUNCTION, _progcb);
        curl_easy_setopt(dt->handle, CURLOPT_XFERINFODATA, dt);
        curl_easy_setopt(dt->handle, CURLOPT_NOPROGRESS, 0);
    } else {
        curl_easy_setopt(dt->handle, CURLOPT_XFERINFOFUNCTION, nullptr);
        curl_easy_setopt(dt->handle, CURLOPT_XFERINFODATA, nullptr);
        curl_easy_setopt(dt->handle, CURLOPT_NOPROGRESS, 1);
    }
    dt->progressCallback = progressCallback;
}

bool Downloader::getInfo(DownloadTask *dt, DownloadInfo &di) {
    di.filetime = dt->filetime;
    di.curr = dt->curr;
    di.total = dt->total;
    di.state = dt->state;
    di.url = dt->url;
    return true;
}

bool Downloader::getContent(DownloadTask *dt, std::string &content) {
    if (dt->state != Finished || !dt->ofile) return false;
    std::ostringstream *ostr = dynamic_cast<std::ostringstream*>(dt->ofile);
    if (ostr == nullptr) return false;
    content = ostr->str();
    return true;
}

DownloadTask *Downloader::addCommon(const std::string &url, bool paused, uint64_t resume) {
    CURL *handle = curl_easy_init();
    if (handle == nullptr) return nullptr;
    DownloadTask *dt = new DownloadTask;
    dt->host = this;
    dt->handle = handle;
    dt->url = url;
    dt->curr = 0ULL;
    dt->total = 0ULL;
    addCommon(dt, paused, resume);
    int uid = ++uid_;
    dt->id = uid;
    tasks_[uid] = dt;
    return dt;
}

void Downloader::addCommon(DownloadTask *dt, bool paused, uint64_t resume) {
    curl_easy_setopt(dt->handle, CURLOPT_PRIVATE, dt);
    curl_easy_setopt(dt->handle, CURLOPT_URL, dt->url.c_str());
    curl_easy_setopt(dt->handle, CURLOPT_FOLLOWLOCATION, 1);
    curl_easy_setopt(dt->handle, CURLOPT_FILETIME, 1);
    curl_easy_setopt(dt->handle, CURLOPT_NOPROGRESS, 1);
    if (!proxy_.empty())
        curl_easy_setopt(dt->handle, CURLOPT_PROXY, proxy_.c_str());
    curl_easy_setopt(dt->handle, CURLOPT_VERBOSE, 1);
    if (resume) {
        curl_easy_setopt(dt->handle, CURLOPT_RESUME_FROM_LARGE, resume);
        if (dt->progressCallback) {
            curl_easy_setopt(dt->handle, CURLOPT_XFERINFOFUNCTION, _progcb);
            curl_easy_setopt(dt->handle, CURLOPT_XFERINFODATA, dt);
            curl_easy_setopt(dt->handle, CURLOPT_NOPROGRESS, 0);
        } else {
            curl_easy_setopt(dt->handle, CURLOPT_XFERINFOFUNCTION, nullptr);
            curl_easy_setopt(dt->handle, CURLOPT_XFERINFODATA, nullptr);
            curl_easy_setopt(dt->handle, CURLOPT_NOPROGRESS, 1);
        }
        if (dt->noBody) {
            curl_easy_setopt(dt->handle, CURLOPT_NOBODY, 1);
        } else {
            curl_easy_setopt(dt->handle, CURLOPT_WRITEFUNCTION, _writecb);
            curl_easy_setopt(dt->handle, CURLOPT_WRITEDATA, dt);
        }
    }
    if (paused) {
        curl_easy_pause(dt->handle, CURLPAUSE_ALL);
        dt->state = Paused;
    } else
        dt->state = Downloading;
    curl_multi_add_handle(handle_, dt->handle);
}

size_t Downloader::_writecb(void *ptr, size_t size, size_t nmemb, void *data) {
    size_t sz = size * nmemb;
    auto *dt = ((DownloadTask*)data);
    dt->host->writeCb(dt, ptr, sz);
    return sz;
}

int Downloader::_progcb(void *data, uint64_t dltotal, uint64_t dlnow, uint64_t ultotal, uint64_t ulnow) {
    auto *dt = ((DownloadTask*)data);
    dt->host->progressCb(dt, dlnow, dltotal);
    return 0;
}

void Downloader::writeCb(DownloadTask *dt, const void *ptr, size_t sz) {
    dt->ofile->write((const char*)ptr, sz);
}

void Downloader::progressCb(DownloadTask *dt, uint64_t curr, uint64_t sz) {
    if (curr == dt->curr && sz == dt->total) return;
    dt->curr = curr;
    dt->total = sz;
    if (dt->progressCallback) dt->progressCallback(dt, curr, sz);
}
